import numpy as np
import matplotlib.pyplot as plt
import os


# Hyper Parameters
hyper_params = {
    'max_x': 8,
    'max_y': 11
}


class RSSIInfo():
    def __init__(self, t, SSID, MAC, RSSI, X, Y):
        self.t = t
        self.SSID = SSID
        self.MAC = MAC
        self.RSSI = RSSI
        self.X = X
        self.Y = Y


class RSSIList():
    def __init__(self, rssi_min_value=-100, map_data='map.txt'):
        self.data_list: list[RSSIInfo] = []
        self.rssi_min_value = rssi_min_value
        self.max_x = hyper_params['max_x']
        self.max_y = hyper_params['max_y']
        self.map_data = map_data
        self.mac_map = {}
        self.pos_to_list = [[[] for _ in range(self.max_y)] for _ in range(self.max_x)]
        self.mac_cnt = 0
        self.device_cnt = 0
        
        self.init_map()

    def init_map(self):
        with open(self.map_data, 'rt') as f:
            self.device_cnt = int(f.readline().strip())
            for line in f.readlines():
                mac, i = line.strip().split(',')
                self.mac_map[mac] = int(i)

    def append(self, rssi_info: RSSIInfo):
        self.data_list.append(rssi_info)
        self.pos_to_list[rssi_info.X][rssi_info.Y].append(rssi_info)

    def get_mean_rssi_at_mac_pos(self, MAC, X, Y, value_fix=True):
        found_cnt = 0
        total_rssi = 0
        for wifi in self.data_list:
            if wifi.MAC == MAC and wifi.X == X and wifi.Y == Y:
                total_rssi += wifi.RSSI
                found_cnt += 1

        if value_fix and found_cnt == 0:
            dx = [1, 0, -1, 0]
            dy = [0, -1, 0, 1]

            for i in range(4):
                if (near_rssi := self.get_mean_rssi_at_mac_pos(MAC, X+dx[i], Y+dy[i], value_fix=False)) > self.rssi_min_value:
                    found_cnt += 1
                    total_rssi += near_rssi

        return total_rssi / found_cnt if found_cnt > 0 else self.rssi_min_value

    def get_list_at_pos(self, X, Y):
        return self.pos_to_list[X][Y]

    def get_list_at_mac(self, MAC):
        ans = []
        for wifi in self.data_list:
            if wifi.MAC == MAC:
                ans.append(wifi)

        return ans

    def get_mac_list(self):
        ans = []
        for wifi in self.data_list:
            if wifi.MAC not in ans:
                ans.append(wifi.MAC)

        return ans

def walkFile(file):
    for root, dirs, files in os.walk(file):
        # 遍历文件
        for f in files:
            yield os.path.join(root, f)


def init_list(file_path: str):
    data_list = RSSIList()
    for file in walkFile(file_path):
        f = open(file, 'rt', encoding='utf8')
        for line in f.readlines():
            t, SSID, MAC, RSSI, X, Y = line.rstrip('\n').split('::')
            data_list.append(RSSIInfo(t, SSID, MAC,
                                    int(RSSI), int(X), int(Y)))

    return data_list

if __name__ == '__main__':
    rssi_list = init_list(file_path='data')
    print(rssi_list.mac_cnt)

    mac_list = rssi_list.get_mac_list()
    print(len(mac_list))
    ans_list = []

    for mac in mac_list:
        ans_list.append((mac, len(rssi_list.get_list_at_mac(mac))))

    ans_list = sorted(ans_list, key=lambda x: x[1], reverse=True)
    print(ans_list)

    # Draw Heat Map
    for MAC_ADDRESS, _ in ans_list:
        # MAC_ADDRESS = '6c:f3:7f:5a:e4:70'
        heat_map_array: np.ndarray = np.zeros(
            [hyper_params['max_y'], hyper_params['max_x']])

        for i in range(hyper_params['max_y']):
            for j in range(hyper_params['max_x']):
                print('%.1f' % rssi_list.get_mean_rssi_at_mac_pos(
                    MAC_ADDRESS, j, i), end=' ')
                heat_map_array[i, j] = rssi_list.get_mean_rssi_at_mac_pos(
                    MAC_ADDRESS, j, i)
            print()

        ax = plt.gca()
        plt.imshow(heat_map_array, interpolation='nearest', cmap=plt.cm.hot,)
        plt.xlabel('X')
        plt.ylabel('Y')
        plt.title(f'{MAC_ADDRESS=}')
        plt.colorbar()
        ax.yaxis.set_ticks_position('left')
        ax.invert_yaxis()
        plt.show()
        plt.close()
